#include <stdio.h>
#include <time.h>
#include <stdlib.h>
#include "gemm.h"

#define A( i, j )     A[ (j)*lda + (i) ]
#define B( i, j )     B[ (j)*ldb + (i) ]
#define C( i, j )     C[ (j)*ldc + (i) ]
#define C_ref( i, j ) C_ref[ (j)*ldc_ref + (i) ]

void test_gemm(int m, int n, int k)
{
    float *A, *B, *C, *C_ref;
    int i, j, p, nx;
    int lda, ldb, ldc, ldc_ref;

    printf(" m %d n %d k %d\n", m, n, k);

    A = (float*)malloc(sizeof(float) * m * k);
    B = (float*)malloc(sizeof(float) * k * n);

    lda = m;
    ldb = k;
    ldc = n;
    ldc_ref = m;

    C = (float*)malloc(sizeof(float) * m * n);
    C_ref = (float*)malloc(sizeof(float) * m * n);
    
    // srand48(time(NULL));
    printf("A: \n");
    for (p = 0; p < k; p++) {
        for (i = 0; i < m; i++) {
            A[p * lda + i] = (float)drand48();
            A[p * lda + i] = 0.1 * (p * lda + i);
            printf("%0.3f ", A[p * lda + i]);
        }
        printf("\n");
    }

    printf("B: \n");
    for (j = 0; j < n; j++) {
        for (p = 0; p < k; p++) {
            B[j * ldb + p] = (float)drand48();
            B[j * ldb + p] = 0.5;
            printf("%0.3f ", B[j * ldb + p]);
        }
        printf("\n");
    }

    for (j = 0; j < n; j++) {
        for (i = 0; i < m; i++) {
            C[j * ldc + i] = (float)0.0;
            C_ref[j * ldc_ref + i] = (float)0.0;
        }
    }

    printf("C: \n");
    gemm_cpu(1, 1, m, n, k, 1.0, A, lda, B, ldb, 1.0, C, ldc);
    for (i = 0; i < m; i++) {
        for (j = 0; j < n; j++) {
            printf("%0.3f ", C[i * ldc + j]);
        }
        printf("\n");
    }

    printf("C_ref: \n");
    for (i = 0; i < m; i++) {
        for (j = 0; j < n; j++) {
            for (p = 0; p < k; p++) {
                C_ref[j * ldc_ref + i] += A[p * lda + i] * B[j * ldb + p];
            }
            printf("%0.3f ", C_ref[j * ldc_ref + i]);
        }
        printf("\n");
    }


    if (A) {
        free(A);
    }

    if (B) {
        free(B);
    }

    if (C) {
        free(C);
    }

    if (C_ref) {
        free(C_ref);
    }
}
int main()
{
    int ret = 0;

    printf("hello ds.\n");
    test_gemm(3, 2, 4);

    return ret;
}